/*
 * nmf_main.cpp
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * main file for maxtrix factorization program
 * Example run (use --help for help on parameters)
 *  Generate data
 *   mpirun -np 16 ./nmf  --gen --data=m1000_10_200_16 --m=200 --n=200 --k=10 
 *  Sequential NMF (ALS)
 *   mpirun -np 1 ./nmf  --anls --m=200 --n=200 --k=10 --iter=1000 --rho=10
 *  Master/Slave NMF
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --iter=5000 --interval=1000
 *  Decentralized NMF (specify network)
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --graph=grid --iter=5000 --interval=1000
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --graph=equal --degree=2 --iter=5000 --interval=1000
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#define ARMA_USE_LAPACK

#include <iostream>
#include <armadillo>

#include "helper/helper.h"
#include "helper/StopWatch.h"
#include <boost/mpi.hpp>
boost::mpi::communicator* p_world = NULL; // p_world is intentionaly placed before
#include "helper/mpi_helper.h"                   // including mpi_helper.h

#include <opencv2/opencv.hpp>
#include <iomanip>

po::variables_map vm;
po::options_description desc("Allowed options");

// JOBS
void run_anls();
void run_master_slave_anls();
void run_dist_anls();
void run_dist_anls_less();
void generate_data();
void convert_image();
void run_bcd();
void run_bpp();
void run_map_reduce_bpp();
void run_map_reduce_bpp_manyrhs();

int main(int argc, char** argv)
{
  mpi::environment env(argc, argv);
  mpi::communicator world;
  p_world = &world;
  
  //std::cerr << "rank=" << MPI_RANK << std::endl;

  desc.add_options()
    ("help", "produce help message")
    ("verbose", "print all debug information")
    ("interval", po::value<int>()->default_value(100), "reporting interval")
    ("seed", po::value<int>()->default_value(120), "random seed")
    // DATA
    ("data", po::value<std::string>(), "data file")
    ("output", po::value<std::string>(), "output basis file")
    ("iterinfo", po::value<std::string>(), "iteration information file")
    // MODEL
    ("model", po::value<std::string>(), "output model file")
    // METHODS
    ("gen", "generate random data (MPI_SIZE)\n--m --n --k")
    ("anls-seq", "solve A = WH^T by Sequential Alternating Nonnegative Least Square\n--data --rate --rho --iter --k")
    ("anls-dist", "solve A = WH^T by Distributed Alternating Nonnegative Least Square\n--data --rate --rho --graph [--degree] --iter --k")
    ("lf-dist", "solve A = UV^T by Distributed Low-Rank Matrix Factorization\n--data --alpha --rho --graph [--degree] --iter --k")
    ("image", "read images")
    ("bcd-seq", "Block coordinate descent (Sequential)")
    ("bpp-seq", "Block principal pivoting (Sequential)")
    ("bpp-dist", "Block principal pivoting (Sequential)")
    ("bpp-dist-rhs", "Block principal pivoting (Sequential) with many RHS")
    // COMMON PARAMETERS
    ("iter", po::value<int>()->default_value(1000), "maximum iterations")
    ("rate", po::value<double>()->default_value(1.0), "learning rate / step size")
    ("density", po::value<double>()->default_value(0.3), "sparsity density")
    ("rho", po::value<double>()->default_value(0.1), "augmented penalty")
    ("rho-incr", po::value<double>()->default_value(2), "penalty increase factor")
    ("rho-decr", po::value<double>()->default_value(0.5), "penalty decrease factor")
    ("alpha", po::value<double>()->default_value(0.1), "concensus penalty")
    ("normalize", "keep columns of H normalized (norm <= 1)")
    ("l1", po::value<double>(), "l1 regularization")
    ("l2", po::value<double>(), "l2 regularization")
    // DISTRIBUTED ALGORITHMS
    ("graph", po::value<std::string>(), "decentralized ADMM with bipartie graph (grid|equal)")
    ("degree", po::value<int>(), "decentralized ADMM with bipartie graph")
    ("less", "use less variables version")
    // GENERATE DATA
    ("m", po::value<int>()->default_value(10), "= A.nrows")
    ("n", po::value<int>()->default_value(10), "= A.ncols")
    ("k", po::value<int>()->default_value(2), "= W.ncols = H.nrows")
    ("np", po::value<int>()->default_value(2), "Number of processes to generate data")
    ("sparse", po::value<double>()->default_value(1), "sparsity level")
    ;

  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm); 

  PO_VARIABLE(seed, vm, "seed", int);
  srand(seed+MPI_RANK);

  if (vm.count("help")) {
    std::cerr << desc << std::endl;
  }
  else if (vm.count("anls-seq")) {
    run_anls();
  }
  else if (vm.count("anls-dist")) {
    if (vm.count("graph")) { // graph --> Decentralized ALS
      if (vm.count("less"))
	run_dist_anls_less();
      else
	run_dist_anls();
    }
    else
      run_master_slave_anls(); // no graph --> Master/Slave ALS
  }
  else if (vm.count("gen")) { // generate data
    generate_data();
  }
  else if (vm.count("image")) {
    convert_image();
  }
  else if (vm.count("bcd-seq")) {
    run_bcd();
  }
  else if (vm.count("bpp-seq")) {
    run_bpp();
  }
  else if (vm.count("bpp-dist")) {
    run_map_reduce_bpp();
  }
  else if (vm.count("bpp-dist-rhs")) {
    run_map_reduce_bpp_manyrhs();
  }
  else {
    MSG(desc);
  }

  return 0;
}

void make_sparse(arma::mat& A, double sparsity)
{
  for (unsigned int i = 0; i < A.n_rows; i++)
    for (unsigned int j = 0; j < A.n_cols; j++)
      if ( arma::randu() > sparsity ) A(i,j) = 0;
}

void generate_data()
{
  PO_REQUIRE(vm, "data");

  PO_VARIABLE(m, vm, "m", int);
  PO_VARIABLE(n, vm, "n", int);
  PO_VARIABLE(k, vm, "k", int);
  PO_VARIABLE(data, vm, "data", std::string);
  PO_VARIABLE(sparsity, vm, "sparse", double);
  
  using namespace arma;
  srand(time(NULL)+MPI_RANK);
  mat W = randu<mat>(MPI_RANK < MPI_SIZE-1 ? m/MPI_SIZE : m - (MPI_SIZE-1)*(m/MPI_SIZE), k);
  make_sparse(W, sparsity);
  mat H;
  if (MPI_RANK == 0) {
    H = randu<mat>(n, k);
    make_sparse(H, sparsity);
  }
  mpi::broadcast(*p_world, H, 0);

  mat A = W*trans(H);
  std::stringstream ss;
  ss << data << "_" << MPI_RANK;
  A.save(ss.str(), arma_binary);
}

void convert_image()
{
  PO_VARIABLE(filename, vm, "data", std::string);
  int m = 2429;
  int n = 19*19;    
  arma::mat A(m, n);
  for (int r = 1; r <= 2429; r++) {
    std::stringstream ss;
    ss << filename << std::setw(4) << std::setfill('0') << r << ".pgm";
    MSG_VERBOSE(vm,ss.str());
    cv::Mat img = cv::imread(ss.str(), 1), gray;
    //MSG(img.rows << " " << img.cols << " " << img.type()
    //  << " " << img.depth() << " " << img.channels());
    cv::cvtColor(img, gray, CV_BGR2GRAY);
    //MSG(gray.rows << " " << gray.cols << " " << gray.type()
    //  << " " << gray.depth() << " " << gray.channels() << " " << CV_8UC1);
    cv::Mat_<unsigned char>& g = (cv::Mat_<unsigned char>&) gray;
    int k=0;
    for(int j = 0; j < g.cols; j++)
      for(int i = 0; i < g.rows; i++,k++) A(r-1,k)=((double)g(i,j)) / 255;
  }
  A.save(filename, arma::arma_binary);
  MSG_VERBOSE(vm,A);
    
  PO_VARIABLE(np, vm, "np", int);
  int start = 0;
  for (int i = 0; i < np; i++, start += m/np) {
    arma::mat B = A.rows( start, i < np-1 ? start + (m/np) - 1 : m-1 );
    std::stringstream ss;
    ss << filename << "_" << np << "_" << i;
    B.save(ss.str(), arma::arma_binary);
  }
}

